#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Neo4j决策树查询工具
功能：查询基于->符号建立的决策树结构
"""

from neo4j import GraphDatabase
import jieba
import re
import sys


class DecisionTreeQuery:
    def __init__(self, uri="bolt://192.168.1.169:7687", username="neo4j", password="Pa@ssw0rd"):
        try:
            self.driver = GraphDatabase.driver(uri, auth=(username, password))
            print(f"✅ 成功连接到Neo4j数据库: {uri}")
        except Exception as e:
            print(f"❌ 连接Neo4j失败: {e}")
            self.driver = None
    
    def close(self):
        if self.driver:
            self.driver.close()
    
    def query_decision_paths(self, product_code):
        """查询产品的所有决策路径"""
        if not self.driver:
            return []
        
        with self.driver.session() as session:
            cypher = """
                MATCH path = (p:Product {code: $product})-[:HAS_DECISION_TREE]->(start:DecisionNode)-[:LEADS_TO*0..10]->(end:DecisionNode)
                RETURN path, start.filename as filename, start.line_number as line_number
                ORDER BY start.filename, start.line_number
            """
            
            print(f"📝 查询决策路径:")
            print(f"   {cypher.replace('$product', f'"{product_code}"')}")
            
            try:
                result = session.run(cypher, product=product_code)
                paths = []
                
                for record in result:
                    path = record['path']
                    nodes = []
                    
                    for node in path.nodes:
                        if hasattr(node, 'labels') and 'DecisionNode' in node.labels:
                            nodes.append({
                                'content': node.get('content', ''),
                                'type': node.get('type', ''),
                                'level': node.get('level', 0)
                            })
                    
                    if nodes:
                        paths.append({
                            'nodes': nodes,
                            'filename': record['filename'],
                            'line_number': record['line_number']
                        })
                
                return paths
                
            except Exception as e:
                print(f"❌ 查询失败: {e}")
                return []
    
    def query_by_keyword(self, keyword):
        """根据关键词查询决策路径"""
        if not self.driver:
            return []
        
        with self.driver.session() as session:
            cypher = """
                MATCH (n:DecisionNode)
                WHERE n.content CONTAINS $keyword
                MATCH path = (start:DecisionNode)-[:LEADS_TO*0..10]->(n)-[:LEADS_TO*0..10]->(end:DecisionNode)
                WHERE start.level = 0
                RETURN path, n.filename as filename, n.line_number as line_number, n.content as matched_content
                ORDER BY n.filename, n.line_number
                LIMIT 10
            """
            
            print(f"📝 关键词查询:")
            print(f"   {cypher.replace('$keyword', f'"{keyword}"')}")
            
            try:
                result = session.run(cypher, keyword=keyword)
                paths = []
                
                for record in result:
                    path = record['path']
                    nodes = []
                    
                    for node in path.nodes:
                        if hasattr(node, 'labels') and 'DecisionNode' in node.labels:
                            nodes.append({
                                'content': node.get('content', ''),
                                'type': node.get('type', ''),
                                'level': node.get('level', 0)
                            })
                    
                    if nodes:
                        # 按level排序
                        nodes.sort(key=lambda x: x['level'])
                        paths.append({
                            'nodes': nodes,
                            'filename': record['filename'],
                            'line_number': record['line_number'],
                            'matched_content': record['matched_content']
                        })
                
                return paths
                
            except Exception as e:
                print(f"❌ 查询失败: {e}")
                return []
    
    def query_solutions_for_problem(self, problem_keyword):
        """查询特定问题的解决方案"""
        if not self.driver:
            return []
        
        with self.driver.session() as session:
            cypher = """
                MATCH (problem:DecisionNode)
                WHERE problem.content CONTAINS $keyword AND problem.type IN ['problem', 'category']
                MATCH path = (problem)-[:LEADS_TO*1..10]->(solution:DecisionNode)
                WHERE solution.type = 'solution'
                RETURN problem.content as problem, solution.content as solution, 
                       problem.filename as filename, problem.line_number as line_number,
                       problem.product_code as product
                ORDER BY problem.filename, problem.line_number
                LIMIT 10
            """
            
            print(f"📝 问题解决方案查询:")
            print(f"   {cypher.replace('$keyword', f'"{problem_keyword}"')}")
            
            try:
                result = session.run(cypher, keyword=problem_keyword)
                solutions = []
                
                for record in result:
                    solutions.append({
                        'problem': record['problem'],
                        'solution': record['solution'],
                        'filename': record['filename'],
                        'line_number': record['line_number'],
                        'product': record['product']
                    })
                
                return solutions
                
            except Exception as e:
                print(f"❌ 查询失败: {e}")
                return []
    
    def get_complete_decision_chain(self, filename, line_number):
        """获取完整的决策链"""
        if not self.driver:
            return []
        
        with self.driver.session() as session:
            cypher = """
                MATCH (start:DecisionNode {filename: $filename, line_number: $line_number, level: 0})
                MATCH path = (start)-[:LEADS_TO*0..10]->(end:DecisionNode)
                WHERE end.filename = $filename AND end.line_number = $line_number
                RETURN nodes(path) as chain
                ORDER BY length(path) DESC
                LIMIT 1
            """
            
            print(f"📝 完整决策链查询:")
            print(f"   {cypher.replace('$filename', f'"{filename}"').replace('$line_number', str(line_number))}")
            
            try:
                result = session.run(cypher, filename=filename, line_number=line_number)
                record = result.single()
                
                if record:
                    chain = []
                    for node in record['chain']:
                        chain.append({
                            'content': node.get('content', ''),
                            'type': node.get('type', ''),
                            'level': node.get('level', 0)
                        })
                    
                    # 按level排序
                    chain.sort(key=lambda x: x['level'])
                    return chain
                
                return []
                
            except Exception as e:
                print(f"❌ 查询失败: {e}")
                return []
    
    def display_decision_paths(self, paths, title="决策路径"):
        """显示决策路径"""
        if not paths:
            print(f"\n❌ 没有找到{title}")
            return
        
        print(f"\n🎯 {title} (共{len(paths)}条):")
        print("=" * 80)
        
        for i, path in enumerate(paths, 1):
            print(f"\n--- 路径 {i} ---")
            if 'filename' in path:
                print(f"📁 文件: {path['filename']} 行{path['line_number']}")
            
            if 'matched_content' in path:
                print(f"🔍 匹配内容: {path['matched_content']}")
            
            if 'problem' in path and 'solution' in path:
                print(f"❓ 问题: {path['problem']}")
                print(f"💡 解决方案: {path['solution']}")
                if 'product' in path:
                    print(f"📦 产品: {path['product']}")
            else:
                print("🔗 决策链:")
                for j, node in enumerate(path['nodes']):
                    arrow = " -> " if j > 0 else ""
                    type_icon = self._get_type_icon(node['type'])
                    print(f"   {arrow}{type_icon} {node['content']}")
    
    def display_decision_chain(self, chain, title="完整决策链"):
        """显示完整决策链"""
        if not chain:
            print(f"\n❌ 没有找到{title}")
            return
        
        print(f"\n🔗 {title}:")
        print("=" * 60)
        
        for i, node in enumerate(chain):
            arrow = " -> " if i > 0 else ""
            type_icon = self._get_type_icon(node['type'])
            print(f"{arrow}{type_icon} {node['content']}")
    
    def _get_type_icon(self, node_type):
        """获取节点类型图标"""
        icons = {
            'product': '📦',
            'problem': '❓',
            'category': '📂',
            'question': '❔',
            'answer': '✅',
            'condition': '🔄',
            'solution': '💡',
            'step': '📝'
        }
        return icons.get(node_type, '🔸')
    
    def smart_query(self, query_text):
        """智能查询"""
        print(f"🔍 智能查询: '{query_text}'")
        print("=" * 60)
        
        # 提取产品编号
        product_codes = re.findall(r'([A-Z]{1,3}\d{3,4})', query_text.upper())
        
        # 提取关键词
        words = list(jieba.cut(query_text))
        keywords = [w.strip() for w in words if len(w.strip()) > 1 and w.strip() not in ['的', '是', '有', '和', '怎么', '什么']]
        
        print(f"🎯 识别的产品编号: {product_codes}")
        print(f"🔍 提取的关键词: {keywords}")
        
        # 策略1: 如果有产品编号，查询该产品的决策路径
        if product_codes:
            for product_code in product_codes:
                print(f"\n📦 查询产品 {product_code} 的决策路径:")
                paths = self.query_decision_paths(product_code)
                self.display_decision_paths(paths, f"{product_code} 决策路径")
        
        # 策略2: 根据关键词查询
        for keyword in keywords[:3]:  # 限制关键词数量
            if len(keyword) > 1 and keyword.upper() not in product_codes:
                print(f"\n🔍 关键词查询: {keyword}")
                paths = self.query_by_keyword(keyword)
                self.display_decision_paths(paths, f"包含'{keyword}'的决策路径")
                
                # 查询解决方案
                solutions = self.query_solutions_for_problem(keyword)
                self.display_decision_paths(solutions, f"'{keyword}'问题的解决方案")


def main():
    """主函数"""
    if len(sys.argv) > 1:
        # 命令行模式
        query = ' '.join(sys.argv[1:])
        
        query_tool = DecisionTreeQuery()
        if not query_tool.driver:
            return
        
        try:
            query_tool.smart_query(query)
        finally:
            query_tool.close()
    
    else:
        # 交互模式
        print("🔍 Neo4j决策树查询工具")
        print("=" * 40)
        print("💡 输入产品编号查看决策路径")
        print("💡 输入关键词搜索相关决策")
        print("💡 输入 'quit' 退出")
        print("=" * 40)
        
        query_tool = DecisionTreeQuery()
        if not query_tool.driver:
            return
        
        try:
            while True:
                query = input("\n🔍 请输入查询: ").strip()
                
                if query.lower() in ['quit', 'exit', '退出']:
                    print("👋 再见!")
                    break
                
                if not query:
                    continue
                
                query_tool.smart_query(query)
        
        finally:
            query_tool.close()


if __name__ == "__main__":
    main()
